"""
Gradio Web App: ensemble model for Cat vs Dog

行为:
- 启动时尝试加载序列化模型文件 `ensemble_model.pkl` (若存在)
- 若没有已保存模型, 会尝试查找特征缓存 `X_*.pkl` / `y_*.pkl` 并训练一个简单分类器并保存
- 支持两类特征方案:
  - vgg: 使用 VGG16 (include_top=False, pooling='max') 提取 512-d 特征
  - flat: 使用灰度 32x32 展平像素
- 提供 Gradio 上传图片界面, 返回预测标签及置信度

注意:
- 运行前请在 PowerShell 中安装依赖: pip install gradio scikit-learn opencv-python tensorflow (视是否需要 vgg)
"""

import os
import glob
import pickle
import logging
from typing import Tuple, Dict, Any

import numpy as np
import cv2

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def find_feature_files() -> Tuple[str, str]:
    """在工作目录查找 X_*.pkl / y_*.pkl 或 X.pkl / y.pkl
    返回 (x_path, y_path) 或 (None, None)
    """
    # 优先 vgg 再 flat 再 通用 X.pkl
    candidates = glob.glob('X_*.pkl') + glob.glob('X.pkl')
    for xp in candidates:
        # 对应的 y 文件
        base = os.path.splitext(xp)[0]
        yp = base.replace('X', 'y') + '.pkl' if not base.endswith('_') else base + 'y.pkl'
        # simpler: try common patterns
        alt_yp = xp.replace('X', 'y', 1)
        if os.path.exists(alt_yp):
            return xp, alt_yp
        if os.path.exists(yp):
            return xp, yp
    return None, None


def infer_label_map_from_dirs() -> Dict[int, str]:
    """尝试通过查找常见数据目录(例如 ./data/train, ./train)来恢复类别名映射（按字母排序）
    返回 idx->name 映射，找不到则返回 {}
    """
    search_dirs = ['data/train', 'train', 'data', './data/train', './train', '.']
    for d in search_dirs:
        if not os.path.exists(d) or not os.path.isdir(d):
            continue
        # 找到直接包含类子目录且子目录包含图片的目录
        subs = [os.path.join(d, s) for s in os.listdir(d) if os.path.isdir(os.path.join(d, s))]
        labels = []
        for s in subs:
            # 检查是否包含图片
            for ext in ('*.jpg', '*.jpeg', '*.png', '*.bmp'):
                if glob.glob(os.path.join(s, ext)):
                    labels.append(os.path.basename(s))
                    break
        if labels:
            labels = sorted(labels)
            return {i: name for i, name in enumerate(labels)}
    return {}


def load_or_train_model() -> Dict[str, Any]:
    """加载已保存模型字典，若不存在则尝试用 X_*.pkl/y_*.pkl 训练并保存。
    返回包含 keys: pipeline, label_map (idx->name), feature_method
    """
    model_file = 'ensemble_model.pkl'
    if os.path.exists(model_file):
        logging.info(f"加载已保存模型: {model_file}")
        with open(model_file, 'rb') as f:
            data = pickle.load(f)
        return data

    # 未发现已保存模型 -> 查找特征缓存
    xp, yp = find_feature_files()
    if xp is None or yp is None:
        logging.error('未找到已保存模型或特征缓存 (X_*.pkl / y_*.pkl)。请先训练并保存模型，或把缓存放到当前目录。')
        return None

    logging.info(f'发现特征缓存: {xp}, {yp}，正在加载并训练分类器...')
    with open(xp, 'rb') as fx:
        X = pickle.load(fx)
    with open(yp, 'rb') as fy:
        y = pickle.load(fy)

    X = np.array(X)
    y = np.array(y)

    # 推断 feature 方法
    feature_method = 'flat'
    if 'vgg' in os.path.basename(xp).lower() or (X.ndim == 2 and X.shape[1] == 512):
        feature_method = 'vgg'

    # 训练一个轻量分类器（LogisticRegression），若需要可改为 RandomForest
    try:
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import StandardScaler
        from sklearn.linear_model import LogisticRegression
    except Exception as e:
        logging.exception('scikit-learn 未安装或导入失败，请安装 scikit-learn')
        return None

    # 检查标签种类数，避免单类数据导致一些分类器报错
    classes = np.unique(y)
    logging.info(f'开始训练 (样本数={X.shape[0]}, 特征维度={X.shape[1]}, classes={classes})')
    if classes.size == 1:
        # 只有一个类别：使用 DummyClassifier 始终预测该类，避免 LogisticRegression 报错
        try:
            from sklearn.dummy import DummyClassifier
        except Exception:
            logging.exception('无法导入 DummyClassifier（sklearn 可能未安装）')
            return None
        clf = DummyClassifier(strategy='constant', constant=classes[0])
        pipeline = make_pipeline(StandardScaler(), clf)
        pipeline.fit(X, y)
        logging.warning(f'训练集仅包含单一类别 {classes[0]}，已使用 DummyClassifier 作为回退')
    else:
        clf = LogisticRegression(max_iter=2000)
        pipeline = make_pipeline(StandardScaler(), clf)
        pipeline.fit(X, y)

    # 尝试推断类别名称映射
    label_map = infer_label_map_from_dirs()
    if not label_map:
        # fallback: 使用数字名称
        unique = sorted(list(set(int(x) for x in y)))
        label_map = {i: f'class_{i}' for i in unique}

    data = {'pipeline': pipeline, 'label_map': label_map, 'feature_method': feature_method}

    # 保存以便下次快速启用
    try:
        with open(model_file, 'wb') as f:
            pickle.dump(data, f)
        logging.info(f'训练完成并保存模型为 {model_file}')
    except Exception:
        logging.exception('保存模型失败（但训练已完成）')

    return data


def extract_vgg_feature(pil_image) -> np.ndarray:
    """使用 Keras VGG16 提取 512-d 特征；输入为 PIL.Image"""
    try:
        from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
        from tensorflow.keras.preprocessing.image import img_to_array
    except Exception:
        raise RuntimeError('缺少 tensorflow 或 Keras，无法使用 VGG16 特征提取。请安装 tensorflow。')

    # 懒加载模型到模块缓存
    global _vgg_model
    if '_vgg_model' not in globals():
        _vgg_model = VGG16(weights='imagenet', include_top=False, pooling='max')

    # PIL -> RGB numpy
    img = pil_image.convert('RGB')
    img = img.resize((224, 224))
    arr = img_to_array(img)
    arr = np.expand_dims(arr, axis=0)
    arr = preprocess_input(arr)
    feat = _vgg_model.predict(arr, verbose=0)
    return feat.reshape(-1)


def extract_flat_feature(pil_image) -> np.ndarray:
    # 灰度 32x32 展平
    img = pil_image.convert('L')
    img = img.resize((32, 32))
    arr = np.array(img, dtype=np.float32)
    return arr.flatten()


def create_predict_function(model_dict):
    pipeline = model_dict['pipeline']
    label_map = model_dict['label_map']
    method = model_dict['feature_method']

    def predict(pil_image):
        # 处理空输入
        if pil_image is None:
            return {'error': '未上传图片'}

        try:
            if method == 'vgg':
                feat = extract_vgg_feature(pil_image)
            else:
                feat = extract_flat_feature(pil_image)

            feat = np.array(feat, dtype=np.float32).reshape(1, -1)

            # 预测概率/分数
            try:
                probs = pipeline.predict_proba(feat)[0]
                idx = int(np.argmax(probs))
                prob = float(probs[idx])
            except Exception:
                # 部分分类器没有 predict_proba
                scores = pipeline.predict(feat)
                idx = int(scores[0])
                prob = 1.0

            label = label_map.get(idx, str(idx))
            # 返回为 dict 方便 Gradio 的 Label 输出
            return {label: round(prob, 4)}
        except Exception as e:
            logging.exception('预测失败')
            return {'error': str(e)}

    return predict


if __name__ == '__main__':
    # 加载或训练模型
    model_dict = load_or_train_model()
    if model_dict is None:
        print('未能加载或训练模型。请查看日志以了解详细原因。')
        print('如果你已有训练好的模型，请将其序列化为 ensemble_model.pkl (包含 keys: pipeline,label_map,feature_method)')
        raise SystemExit(1)

    # 延迟导入 gradio，以便在没有安装时给出提示
    try:
        import gradio as gr
    except Exception:
        logging.exception('Gradio 未安装，请运行: pip install gradio')
        raise

    predict_fn = create_predict_function(model_dict)

    # 使用 Gradio 界面
    title = 'Ensemble Cat vs Dog Classifier'
    description = '上传一张图片，模型会判断是猫还是狗（需要先在当前目录放好模型或特征缓存以便自动训练）'

    iface = gr.Interface(fn=predict_fn,
                         inputs=gr.Image(type='pil'),
                         outputs=gr.Label(num_top_classes=2),
                         title=title,
                         description=description,
                         allow_flagging='never')

    iface.launch(server_name='0.0.0.0')
